<!DOCTYPE html>
<html>
<head>
  <title>LoRA Metadata Viewer</title>
  <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/tokyo-night-dark.min.css">
  <style>
    .light-theme { --bg:##F4F4F4; --vltheme:#3e3a3e; --ltheme:#dddbd6; --theme:#e2dfd8; --dtheme:#ebe8e1; --vdtheme:#f1f0e4; }
    .dark-theme { --bg:#121212; --vltheme:#c1c5c1; --ltheme:#2C2E33; --theme:#373A40; --dtheme:#25262b; --vdtheme:#1A1B26; }
    :root { --blue:#4dabf7; --red:#dc3545; }
    body { font-family: Arial, sans-serif; background-color:var(--bg); color: var(--vltheme); margin:0; }
    .body-content { margin:10px; }
    #dropArea { width: 100%; border: 2px dashed var(--theme); text-align: center; padding: 50px 0; margin-bottom: 20px; background-color: var(--dtheme); }
    #dropArea.hover { background-color: var(--ltheme); }
    .collapsible { cursor: pointer; user-select: none; margin-bottom: 5px; }
    .collapsible:hover { background-color: var(--vdtheme); }
    .collapsible::before { content: "▼ "; display: inline-block; width: 1em; }
    .active::before { content: "▲ "; }
    .content { display: none; padding: 5px; position: relative; }
    .copy-btn { position: absolute; top: 1.5em; right: 0.5em; background-color: transparent; border: none; color: grey; }
    .copy-btn:hover { background-color: grey; color: white; }
    .note { font-size: 0.7em; font-style: italic; }
    .header { margin:20px; font-size: 2em; font-weight: bold; }
    .footer { margin:10px; text-align: center; font-size: 0.9em; }
    #notificationPanel, #editorNotificationPanel { text-align: center; color:var(--red); font-weight: bold; }
    a { color: var(--blue); font-weight: bold; }
    textarea { width:100%; color: var(--vltheme); background-color: var(--ltheme); }
    hr { border-color: var(--theme) }
    input[type=checkbox] { margin: 10px;}
    .editor-header { text-align: center; margin: 10px; }
    .editor-header button { font-size: 1.2em; color: var(--vltheme); background-color: var(--ltheme); border: none; border-radius: 10px; padding:8px; }
    .editor-header button:hover { background-color: var(--theme); }
  </style>
</head>
<body class="light-theme">
  <div class="header">
    LoRA Metadata Viewer
  </div>
  <hr/>

  <div class="body-content">
    <div id="dropArea" class="drop-area">
      <p>Drag and drop your safetensors file here </br>or click to select a safetensors file</p>
    </div>

    <div id="notificationPanel"></div>

    <h3 class="collapsible">Settings</h3>
    <div class="content">
      <div><input type="checkbox" id="enableSummary" data-display="summaryPanel" checked><label for="enableSummary">Display Summary</label></div>
      <div><input type="checkbox" id="enableCivitAiInfo" data-display="civitAiInfoPanel" checked><label for="enableCivitAiInfo">Display CivitAI Info</label></div>
      <div><input type="checkbox" id="enableTagFrequency" data-display="tagFrequencyPanel" checked><label for="enableTagFrequency">Display Tag Frequency</label></div>
      <div><input type="checkbox" id="enableMetadata" data-display="metadataPanel" checked><label for="enableMetadata">Display Metadata</label></div>
      <div><input type="checkbox" id="enableMetadataEditor" data-display="metadataEditorPanel" checked><label for="enableMetadataEditor">Display Metadata Editor</label></div>
      <div><input type="checkbox" id="enableDateFormat" checked><label for="enableDateFormat">Display date fields in human readable format</label></div>     
      <div><input type="checkbox" id="darkMode" onclick="toggleDarkMode()"><label for="darkMode">Dark Mode</label></div>   
      <hr/>
      <div><b>Summary Field Selection:</b> Specify the name of the fields to be displayed in the summary section, separated by commas:</div>
      <textarea id="summaryFields" rows="8" cols="100">ss_output_name,ss_sd_model_name,ss_network_module,ss_total_batch_size,ss_resolution,ss_optimizer,ss_lr_scheduler,ss_network_module,ss_clip_skip,ss_network_dim,ss_network_alpha,ss_epoch,ss_num_epochs,ss_steps,ss_max_train_steps,ss_learning_rate,ss_text_encoder_lr,ss_unet_lr,ss_shuffle_caption,ss_keep_tokens,ss_flip_aug,ss_noise_offset,ss_adaptive_noise_scale,ss_min_snr_gamma,ss_training_started_at,ss_training_finished_at,sshs_model_hash</textarea>
      <div class="note">Note: If you wish to modify the default list shown here, open this file in a text editor, search for "summaryFields", and modify the values contained within the 'textarea' tag.</div>
    </div>

    <div id="summaryPanel">
      <h3 class="collapsible expanded">Summary</h3>
      <div class="content">
        <button class="copy-btn" onclick="copyToClipboard('summary')">Copy to clipboard</button>
        <pre id="summary"><code class="json">{}</code></pre>
      </div>
    </div>

    <div id="civitAiInfoPanel">
      <h3 class="collapsible">CivitAI Info</h3>
      <div class="content">
        <ul>
          <li><b>URL:</b><a id="modelUrl" target='_blank' href="#"></a></li>
          <li><b>Resource Info:</b><a id="resourceUrl" target='_blank' href="#"></a></li>
          <li><div><b>Preview:</b></div><img id='previewImage' src='' /></li>
        </ul>
        <div><pre id="civitaiDisplay"><code class="json">{}</code></pre></div>
        <div class="note">Note: This section requires an internet connection in order to work. Details will be omited if either the
            resource is not fount in CivitAI or there is no internet connectivity.</div>
      </div>
    </div>

    <div id="tagFrequencyPanel">
      <h3 class="collapsible">Tag Frequency</h3>
      <div class="content">
        <button class="copy-btn" onclick="copyToClipboard('keywordDetails')">Copy to clipboard</button>
        <pre id="keywordDetails"><code class="json">{}</code></pre>
      </div>
    </div>

    <div id="metadataPanel">
      <h3 class="collapsible">Metadata</h3>
      <div class="content">
        <button class="copy-btn" onclick="copyToClipboard('metadataDisplay')">Copy to clipboard</button>
        <pre id="metadataDisplay"><code class="json">{}</code></pre>
      </div>
    </div>

    <div id="metadataEditorPanel">
      <h3 class="collapsible">Metadata Editor</h3>
      <div class="content">
        <div id="editorHeader" class="editor-header" style="display: none;">
          <button onclick="downloadFile()">Update & Download</button>
          <button onclick="downloadFile(true)">Purge & Download</button>
        </div>
        <div id="editorNotificationPanel"></div>
        <textarea id="metadataEditor" rows="100" cols="100"></textarea>
      </div>
    </div>
  </div>

  <hr/>
  <div class="footer">
    LoRA Metadata Viewer v1.4 | 
    <a href="https://github.com/Xypher7/lora-metadata-viewer" target='_blank'>GitHub</a> |
    <a href="https://civitai.com/models/249721/lora-metadata-viewer" target='_blank'>CivitAI</a>
  </div>

  <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
  <script>
    const dropArea = document.getElementById('dropArea');
    const metadataDisplay = document.getElementById('metadataDisplay');
    const metadataEditor = document.getElementById('metadataEditor');
    const summary = document.getElementById('summary');
    const summaryFieldsInput = document.getElementById('summaryFields');
    const modelUrl = document.getElementById('modelUrl');
    const resourceUrl = document.getElementById('resourceUrl');
    const civitaiDisplay = document.getElementById('civitaiDisplay');
    const previewImage = document.getElementById('previewImage');
    const keywordDetails = document.getElementById('keywordDetails');
    const notificationPanel = document.getElementById('notificationPanel');
    const editorNotificationPanel = document.getElementById('editorNotificationPanel');
    const editorHeader = document.getElementById('editorHeader');
    const darkMode = document.getElementById('darkMode');

    const enableSummary = document.getElementById('enableSummary');
    const enableCivitAiInfo = document.getElementById('enableCivitAiInfo');
    const enableTagFrequency = document.getElementById('enableTagFrequency');
    const enableMetadata = document.getElementById('enableMetadata');
    const enableMetadataEditor = document.getElementById('enableMetadataEditor');
    const enableDateFormat = document.getElementById('enableDateFormat');

    let safetensorsFile;

    if(window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches) { 
      toggleDarkMode();
      darkMode.checked = !darkMode.checked;
    }

    try {
      hljs.highlightElement(summary.firstChild);
      hljs.highlightElement(metadataDisplay.firstChild);
      hljs.highlightElement(keywordDetails.firstChild);
      hljs.highlightElement(civitaiDisplay.firstChild);
    } catch (error) { }

    // Prevent default drag behaviors
    ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
      dropArea.addEventListener(eventName, preventDefaults, false);
      document.body.addEventListener(eventName, preventDefaults, false);
    });

    // Highlight drop area when a file is dragged over it
    ['dragenter', 'dragover'].forEach(eventName => {
      dropArea.addEventListener(eventName, highlight, false);
    });

    // Remove highlighting when a file is dragged away from the drop area
    ['dragleave', 'drop'].forEach(eventName => {
      dropArea.addEventListener(eventName, unhighlight, false);
    });

    // Handle dropped files
    dropArea.addEventListener('drop', handleDrop, false);

    // Add hover effect using JavaScript
    dropArea.addEventListener('mouseover', () => {
      dropArea.classList.add('hover');
    });

    dropArea.addEventListener('mouseout', () => {
      dropArea.classList.remove('hover');
    });

    // File input handling
    const fileInput = document.createElement('input');
    fileInput.type = 'file';
    fileInput.style.display = 'none';
    document.body.appendChild(fileInput);

    dropArea.addEventListener('click', () => {
      fileInput.click();
    });

    fileInput.addEventListener('change', () => {
      const file = fileInput.files[0];
      handleFile(file);
    });

    // Collapsible panels
    document.querySelectorAll('.collapsible').forEach((collapsible, i) => {
      collapsible.addEventListener('click', function () {
        this.classList.toggle('active');
        const content = this.nextElementSibling;
        if (content.style.display === 'block') {
          content.style.display = 'none';
        } else {
          content.style.display = 'block';
        }
      });

      if (collapsible.classList.contains('expanded')) collapsible.click();
    });

    // Panel display toggle
    document.querySelectorAll('input[type="checkbox"][data-display]').forEach(function(checkbox) {
      checkbox.addEventListener('click', function() {
        const targetId = this.getAttribute('data-display');
        const targetElement = document.getElementById(targetId);
        if (targetElement) {
          targetElement.style.display = this.checked ? 'block' : 'none';
        }
      });
    });

    function preventDefaults(event) {
      event.preventDefault();
      event.stopPropagation();
    }

    function highlight() {
      dropArea.classList.add('hover');
    }

    function unhighlight() {
      dropArea.classList.remove('hover');
    }

    function handleDrop(event) {
      const dt = event.dataTransfer;
      const file = dt.files[0];
      handleFile(file);
      unhighlight();
    }

    function handleFile(file) {
      clearAll();
      safetensorsFile = file;
      if(!file.name.endsWith('.safetensors')) {
        notificationPanel.innerHTML = 'Please drop a valid .safetensors file.';
        return;
      }

      const reader = new FileReader();
      reader.onload = function (event) {
        try {
          const dataView = new DataView(reader.result);
          const metadataSize = dataView.getUint32(0, true);
          const metadataBytes = new Uint8Array(reader.result.slice(8, 8 + metadataSize));
          const textDecoder = new TextDecoder("utf-8");
          const headerStr = textDecoder.decode(metadataBytes);
          const header = JSON.parse(headerStr);
          const formattedMetadata = header['__metadata__'];        

          if(!formattedMetadata) {
            notificationPanel.innerHTML = 'No metadata found';
            return;
          }
          const metadata = {};
          for (const key in formattedMetadata) {
            if (typeof formattedMetadata[key] === 'string') {
              try {
                metadata[key] = JSON.parse(formattedMetadata[key]);
              } catch (error) {
                metadata[key] = formattedMetadata[key];
              }
            } else {
              metadata[key] = formattedMetadata[key];
            }
          }

          updateMetadataEditor(metadata);

          formatDate(metadata, 'ss_training_started_at');
          formatDate(metadata, 'ss_training_finished_at');

          updateMetadata(metadata);
          updateSummary(metadata);
          updateTagFrequency(metadata);
          updateCivitAiInfo(metadata['sshs_model_hash'], file);
        } catch (error) {
          notificationPanel.innerHTML = 'Error parsing metadata.';
          console.error(error);
        }
      };

      reader.readAsArrayBuffer(file);
    }

    function formatDate(metadata, prop){
      if(enableDateFormat.checked && metadata[prop])  
        metadata[prop] = new Date(metadata[prop] * 1000);     
    }

    function clearAll(){
      notificationPanel.innerHTML = '';
      editorNotificationPanel.innerHTML = '';
      summary.innerHTML = '';
      modelUrl.href = '#';
      modelUrl.innerHTML = '';
      resourceUrl.href = '#';
      resourceUrl.innerHTML = '';
      previewImage.src = '';
      civitaiDisplay.innerHTML = '';
      keywordDetails.innerHTML = '';
      metadataDisplay.innerHTML = '';
      metadataEditor.value = '';
      editorHeader.style.display = "none";
    }

    function updateMetadata(metadata) {
      if(enableMetadata.checked) {
        metadataDisplay.innerHTML = '<code class="json">' + JSON.stringify(metadata, null, 2) + '</code>';
        try { hljs.highlightElement(metadataDisplay.firstChild); } catch (error) { }
      }
    }

    function updateMetadataEditor(metadata) {
      if(enableMetadataEditor.checked) {
        metadataEditor.value = JSON.stringify(metadata, null, 2);
        editorHeader.style.display = "block";
      }
    }

    function updateSummary(metadata) {     
      if(enableSummary.checked){
        const summaryFieldList = summaryFieldsInput.value.split(',').map(value => value.trim());;
        let summaryJson = {};
        summaryFieldList.forEach((item, index) => {
          summaryJson[item] = metadata[item];
        });

        summary.innerHTML = '<code class="json">' + JSON.stringify(summaryJson, null, 2) + '</code>';
        try { hljs.highlightElement(summary.firstChild); } catch (error) { }
      }
    }

    function copyToClipboard(elementId) {
      const textToCopy = document.getElementById(elementId);
      const textarea = document.createElement('textarea');
      textarea.value = textToCopy.textContent;
      document.body.appendChild(textarea);
      textarea.select();
      document.execCommand('copy');
      document.body.removeChild(textarea);
      alert('Text copied to clipboard!');
    }

    function updateCivitAiInfo(hash, file) {
      const baseApiUrl = 'https://civitai.com/api/v1/model-versions/by-hash/';
      const baseModelUrl = 'https://civitai.com/models/';
      let finalHash = hash;

      const fetchWithHash = async (url) => {
        try {
          const response = await fetch(url);
          if (!response.ok) {
            return null;
          }
          return response.json();
        } catch (error) {
          console.log('There was a problem fetching the data:', error);
          return null; // Return null if fetching fails
        }
      };

      const fetchWithFallback = async () => {
        let data = hash ? await fetchWithHash(baseApiUrl + hash) : null;
        if (!data) {
          const calculatedHash = await calculateFileHash(file);
          data = await fetchWithHash(baseApiUrl + calculatedHash);
          finalHash = calculatedHash;
        }
        return data;
      };

      if (enableCivitAiInfo.checked) {
        fetchWithFallback()
          .then(data => {
            if (data) {
              modelUrl.href = baseModelUrl + data.modelId;
              modelUrl.innerHTML = modelUrl.href;
              resourceUrl.href = baseApiUrl + finalHash;
              resourceUrl.innerHTML = baseApiUrl + finalHash;
              civitaiDisplay.innerHTML = '<code class="json">' + JSON.stringify(data, null, 2) + '</code>';
              try { hljs.highlightElement(civitaiDisplay.firstChild); } catch (error) {}
              previewImage.src = data.images[0].url;
            } else {
              console.log('No valid response found.');
            }
          })
          .catch(error => {
            console.log('There was a problem:', error);
          });
      }
    }

    function updateTagFrequency(data){
      try {  
        if(enableTagFrequency.checked && data['ss_tag_frequency']){
          // Extracting and flattening the items from both objects
          const items = Object.values(data['ss_tag_frequency']).flatMap(obj => Object.entries(obj));

          // Create an object to store and sum the values
          const groupedItems = {};
          items.forEach(([key, value]) => {
            const [group, tag] = key.split('_');
            if (!groupedItems[group]) groupedItems[group] = {};
            if (!groupedItems[group][tag]) groupedItems[group][tag] = 0;
            groupedItems[group][tag] += value;
          });

          // Convert the summed items to direct tag-value pairs
          const finalResult = {};
          for (const group in groupedItems) {
            let total = 0;
            for (const tag in groupedItems[group]) {
              total += groupedItems[group][tag];
            }
            finalResult[group] = total;
          }

          // Sort the final result based on the numeric values
          const sortedResult = Object.fromEntries(
            Object.entries(finalResult).sort(([, a], [, b]) => b - a)
          );

          keywordDetails.innerHTML = '<code class="json">' + JSON.stringify(sortedResult, null, 2) + '</code>';
          try { hljs.highlightElement(keywordDetails.firstChild); } catch (error) { }
        }
      } catch (error) { console.error(error); }
    }

    function calculateFileHash(file) {
      return new Promise((resolve, reject) => {
        const reader = new FileReader();
        reader.onload = async () => {
          const buffer = reader.result;
          try {
            const hashBuffer = await crypto.subtle.digest('SHA-256', buffer);
            const hashArray = Array.from(new Uint8Array(hashBuffer));
            const hashHex = hashArray.map(byte => ('00' + byte.toString(16)).slice(-2)).join('');
            resolve(hashHex);
          } catch (err) {
            reject(err);
          }
        };

        reader.onerror = () => {
          reject(new Error('Unable to read the file.'));
        };

        reader.readAsArrayBuffer(file);
      });
    }

    function downloadFile(purge) {
      editorNotificationPanel.innerHTML = '';
      let newMetadata
      try {
        newMetadata = metadataEditor.value && !purge ? JSON.parse(metadataEditor.value) : {};
      } catch (e) {
        editorNotificationPanel.innerHTML = "Error parsing edited metadata. Ensure the metadata is in valid JSON format and try again.";
        return;
      }

      try{
        const reader = new FileReader();
        reader.onload = function (event) {
          const dataView = new DataView(reader.result);
          const metadataSize = dataView.getUint32(0, true);
          const metadataBytes = new Uint8Array(reader.result.slice(8, 8 + metadataSize));
          const textDecoder = new TextDecoder("utf-8");
          const headerStr = textDecoder.decode(metadataBytes);
          let header = JSON.parse(headerStr);
          header['__metadata__'] = newMetadata;

          const newHeaderStr = JSON.stringify(header);
          const textEncoder = new TextEncoder();
          const newHeaderBytes = textEncoder.encode(newHeaderStr);

          const newHeaderSize = newHeaderBytes.length;
          const newFileArrayBuffer = new ArrayBuffer(8 + newHeaderSize + reader.result.byteLength - 8 - metadataSize);
          const newFileDataView = new DataView(newFileArrayBuffer);
          newFileDataView.setUint32(0, newHeaderSize, true);
          new Uint8Array(newFileArrayBuffer, 8, newHeaderSize).set(newHeaderBytes);
          new Uint8Array(newFileArrayBuffer, 8 + newHeaderSize).set(new Uint8Array(reader.result, 8 + metadataSize));
          
          const blob = new Blob([newFileArrayBuffer], { type: "application/octet-stream" });
          const url = URL.createObjectURL(blob);
          const a = document.createElement("a");
          a.href = url;
          a.download = safetensorsFile.name.replace(/(\.[^.]+)$/, `_${purge?"purged":"edited"}$1`);
          document.body.appendChild(a);
          a.click();
          document.body.removeChild(a);
          URL.revokeObjectURL(url);

          alert(`Metadata ${purge?"purged":"updated"} successfully!`);
        };
        reader.readAsArrayBuffer(safetensorsFile);
      } catch (error) {
        editorNotificationPanel.innerHTML = "An error occured while updating the file.";
        console.error(error);
      }
    }

    function toggleDarkMode(){
      document.body.classList.toggle('light-theme');
      document.body.classList.toggle('dark-theme');
    }

  </script>
</body>
</html>
